#include "ahuff.h"
#include <stddef.h>
#include <math.h>
#include <crtdbg.h>
#include "ByteWriter.h"
#include "ByteReader.h"
#include "macros.h"

AHuff::AHuff(void)
{
	Init();
}

AHuff::~AHuff(void)
{
	Dispose();
}

//
// Purpose:
//   Performs all required initialization works
//
void AHuff::Init()
{
	//
	// variable/field inits
	//
	ZERO( m_Nodes )
	m_pOut          = NULL;
	m_NextFreeNode  = 0;

	return;
}

//
// Purpose:
//   If any memory is allocated for the output buffer(s) then it gets deallocated.
//
// Returns:
//   Whether it was needed to free the output buffer or not
//
BOOL AHuff::FreeOutBuff()
{
	if ( m_pOut )
	{
		free( m_pOut );
		m_pOut = NULL;
		return TRUE; // OK
	}
	else
		return FALSE; // ERR
}

//
// Purpose:
//   Performs all required class cleanups
//
void AHuff::Dispose()
{
	FreeOutBuff();

	return;
}

//
// Purpose:
//   Initialize the adaptive Huffman tree before the actual work.
//   The special symbols EndOfStream/Escape and the root nodes
//   are inserted.
//
void AHuff::InitTree()
{
	int i;

	// insert root branch
	m_Nodes[ ROOT_NODE ].child           = ROOT_NODE + 1;
	m_Nodes[ ROOT_NODE ].bIsLeaf         = FALSE;
	m_Nodes[ ROOT_NODE ].count           = 2;
	m_Nodes[ ROOT_NODE ].dad             = -1;

	// insert EndOfStream leaf
	m_Nodes[ ROOT_NODE + 1 ].child       = END_OF_STREAM;
	m_Nodes[ ROOT_NODE + 1 ].bIsLeaf     = TRUE;
	m_Nodes[ ROOT_NODE + 1 ].count       = 1;
	m_Nodes[ ROOT_NODE + 1 ].dad         = ROOT_NODE;
	m_Leafs[ END_OF_STREAM ]             = ROOT_NODE + 1;

	// insert Escape leaf
	m_Nodes[ ROOT_NODE + 2 ].child       = ESCAPE;
	m_Nodes[ ROOT_NODE + 2 ].bIsLeaf     = TRUE;
	m_Nodes[ ROOT_NODE + 2 ].count       = 1;
	m_Nodes[ ROOT_NODE + 2 ].dad         = ROOT_NODE;
	m_Leafs[ ESCAPE ]                    = ROOT_NODE + 2;

	m_NextFreeNode                       = ROOT_NODE + 3;

	for ( i=0; i < END_OF_STREAM; i++ ) // exclude the special symbols - Escape/EndOfStream
		m_Leafs[i] = -1; // mark as unused

	return;
}

//
// Purpose:
//   This routine is called after a special number of bytes were
//   already passed. It
//
void AHuff::RebuildTree()
{
    int        i, j, k;
    UINT       count;

	/*
	 * To start rebuilding the table,  I collect all the leaves of the Huffman
	 * tree and put them in the end of the tree.  While I am doing that, I
	 * scale the counts down by a factor of 2.
	 */
	j = m_NextFreeNode - 1;
	for ( i = j; i >= ROOT_NODE; i-- )
	{
		if ( m_Nodes[ i ].bIsLeaf )
		{
			m_Nodes[ j ] = m_Nodes[ i ];
			m_Nodes[ j ].count = ( m_Nodes[ j ].count + 1 ) / 2;
			j--;
		}
	}

	/*
	 * At this point, j points to the first free node.  I now have all the
	 * leaves defined, and need to start building the higher nodes on the
	 * tree. I will start adding the new internal nodes at j.  Every time
	 * I add a new internal node to the top of the tree, I have to check to
	 * see where it really belongs in the tree.  It might stay at the top,
	 * but there is a good chance I might have to move it back down.  If it
	 * does have to go down, I use the memmove() function to scoot everyone
	 * bigger up by one node.  Note that memmove() may have to be change
	 * to memcpy() on some UNIX systems.  The parameters are unchanged, as
	 * memmove and  memcpy have the same set of parameters.
	 */
	for ( i = m_NextFreeNode - 2; j >= ROOT_NODE; i-=2, j-- )
	{
		k = i + 1;
		m_Nodes[ j ].count = m_Nodes[ i ].count + m_Nodes[ k ].count;
		count = m_Nodes[ j ].count;
		m_Nodes[ j ].bIsLeaf = FALSE; 

		for ( k = j + 1; count < m_Nodes[ k ].count; k++ );

		k--;

		memmove( &m_Nodes[ j ], &m_Nodes[ j + 1 ], (k-j) * sizeof(HT_ITEM) );
		m_Nodes[ k ].count   = count;
		m_Nodes[ k ].child   = i;
		m_Nodes[ k ].bIsLeaf = FALSE;
	}

	/*
	 * The final step in tree reconstruction is to go through and set up
	 * all of the leaf and dad members.  This can be safely done now
	 * that every node is in its final position in the tree.
	 */
    for ( i = m_NextFreeNode - 1; i >= ROOT_NODE; i-- )
	{
        if ( m_Nodes[ i ].bIsLeaf )
		{
            k = m_Nodes[ i ].child;
            m_Leafs[ k ] = i;
        }
		else
		{
            k = m_Nodes[ i ].child;
            m_Nodes[ k ].dad = m_Nodes[ k + 1 ].dad = i;
        }
    }

	return;
}

//
// Purpose:
//   Exchange 2 nodes and take care on the site effects
//
// Args:
//   i/j  - indices into the m_Nodes chain
//
void AHuff::SwapNodes( UINT i, UINT j )
{
    HT_ITEM temp;

	// adjust i-node
    if ( m_Nodes[ i ].bIsLeaf )
        m_Leafs[ m_Nodes[ i ].child ] = j;
    else
	{
        m_Nodes[ m_Nodes[ i ].child ].dad = j;
        m_Nodes[ m_Nodes[ i ].child + 1 ].dad = j;
    }

	// adjust j-node
    if ( m_Nodes[ j ].bIsLeaf )
        m_Leafs[ m_Nodes[ j ].child ] = i;
    else
	{
        m_Nodes[ m_Nodes[ j ].child ].dad = i;
        m_Nodes[ m_Nodes[ j ].child + 1 ].dad = i;
    }

	// exchange the node structures in the tree
    temp               = m_Nodes[ i ];
    m_Nodes[ i ]       = m_Nodes[ j ];
    m_Nodes[ i ].dad   = temp.dad;
    temp.dad           = m_Nodes[ j ].dad;
    m_Nodes[ j ]       = temp; 

	return;
}

//
// Purpose:
//   Insert the leaf for a new symbol into the tree
//
// Remarks:
//   the count of this symbol is firstly set to 0
//
void AHuff::AddNewNode( BYTE symbol )
{
    int lightest_node;
    int new_node;
    int zero_count_node;

    lightest_node     = m_NextFreeNode - 1;
    new_node          = m_NextFreeNode;
    zero_count_node   = m_NextFreeNode + 1;
    m_NextFreeNode   += 2;

    m_Nodes[ new_node ]                    = m_Nodes[ lightest_node ];
    m_Nodes[ new_node ].dad                = lightest_node;
    m_Leafs[ m_Nodes[ new_node ].child ]   = new_node;

    m_Nodes[ lightest_node ].child         = new_node;
    m_Nodes[ lightest_node ].bIsLeaf       = FALSE;

    m_Nodes[ zero_count_node ].child       = symbol;
    m_Nodes[ zero_count_node ].bIsLeaf     = TRUE;
    m_Nodes[ zero_count_node ].count       = 0;
    m_Nodes[ zero_count_node ].dad         = lightest_node;
    m_Leafs[ symbol ]                      = zero_count_node; 
      
	return;
}

//
// Purpose:
//   Increment the count of the specified symbol and adjust
//   the tree
//
void AHuff::UpdateModel( BYTE symbol )
{
	UINT  iCurNode, iNewNode;

	// rebuild the whole tree ?
	if ( m_Nodes[ROOT_NODE].count == MAX_COUNT )
	   RebuildTree();

	// increment + update
	iCurNode = m_Leafs[ symbol ];
	while ( iCurNode != -1 )
	{
		m_Nodes[iCurNode].count++;
		for ( iNewNode = iCurNode; iNewNode > ROOT_NODE; iNewNode-- )
			if ( m_Nodes[ iNewNode-1 ].count >= m_Nodes[ iCurNode ].count )
				break;
		if ( iCurNode != iNewNode )
		{
			SwapNodes( iCurNode, iNewNode );
			iCurNode = iNewNode;
		}
		iCurNode = m_Nodes[ iCurNode ].dad;
	}

	return;
}

//
// Purpose:
//   Compress the specified data block and return ptr/size of the output buffer which
//   will be freed automatically within the class destructor
//
BOOL AHuff::EncodeData( IN PBYTE pbyInput, IN DWORD cbInput, OUT PBYTE &pOut, OUT DWORD &dwOut )
{
	//
	// init
	//
	FreeOutBuff();
	InitTree();
	
	// pre-initialize output variables
	pOut  = NULL;
	dwOut = 0;

	//
	// alloate memory
	//
	m_cbOut  = (DWORD)(cbInput * 1.2);
	m_pOut   = (PBYTE)malloc( m_cbOut );
	if ( !m_pOut )
		return FALSE; // ERR

	//
	// run compression process
	//
	DWORD       cbDone = 0; // meant for the input buffer
	ByteWriter  writer;

	writer.Assign( (void*)m_pOut, m_cbOut );

	// compression loop
	BOOL b;
	while( cbDone != cbInput )
	{
		BYTE byCur = pbyInput[cbDone];

		b = EncodeSymbol( writer, byCur );
		UpdateModel( byCur );

		_ASSERT( b );

		cbDone++;
	}
	// paste EndOfStream
	b = EncodeSymbol( writer, END_OF_STREAM );
	_ASSERT( b );

	writer.FinishWriting();

	// post-initialize output variables
	pOut  = m_pOut;
	dwOut = writer.GetWrittenByteCount();

	return TRUE; // OK
}

//
// Purpose:
//   Decompress the specified data block and return ptr/size of the output buffer which
//   will be freed automatically within the class destructor
//
BOOL AHuff::DecodeData( IN PBYTE pbyInput, IN DWORD cbInput, IN DWORD cbUncompressed, OUT PBYTE &pOut, OUT DWORD &dwOut )
{
	FreeOutBuff();
	InitTree();

	//
	// pre-initialize output variables
	//
	pOut   = NULL;
	dwOut  = 0;

	//
	// allocate memory for output
	//
	m_pOut = (PBYTE)malloc( cbUncompressed );
	if ( !m_pOut )
		return FALSE; // ERR
	
	// prepare decompression
	ByteReader  reader;
	PBYTE       pby = m_pOut;
	reader.Assign( pbyInput, cbInput );

	//
	// decompression loop
	//
	while( TRUE )
	{
		// read symbol
		int symbol = DecodeSymbol( reader );
		if ( symbol == END_OF_STREAM )
			break;

		// stamp byte
		*pby = (BYTE)symbol;

		// update tree
		UpdateModel( symbol );

		// adjust out-ptr
		++pby;
	}

	//
	// post-initialize output variables
	//
	pOut  = m_pOut;
	dwOut = cbUncompressed;

	return TRUE; // OK
}

//
// Purpose:
//   Paste the bit sequence of the specified symbol into the buffer
//   maintained by the passed ByteWriter class example
//
// Remarks:
//   If symbol is stamped the first time into the output buffer then
//   the pure 8bit literal, prefixed by an escape code, will be printed
//
BOOL AHuff::EncodeSymbol( ByteWriter &writer, int symbol )
{
    DWORD  code, iCurBit;
    int    iCodeSize, iCurNode;

    code       = 0;
    iCurBit    = 0x80000000; // highest bit in 4byte cardinal is set
    iCodeSize  = 0;
    iCurNode   = m_Leafs[ symbol ];

	// if printed first time, take an escape code at first
    if ( iCurNode == -1 )
		iCurNode = m_Leafs[ ESCAPE ]; 

    while ( iCurNode != ROOT_NODE )
	{
        if ( (iCurNode & 1) == 0 )
            code |= iCurBit;
        iCurBit >>= 1;
        iCodeSize++;

        iCurNode = m_Nodes[ iCurNode ].dad;
    }

	code >>= 32 - iCodeSize; // move the bits to the low
    BOOL b = writer.WriteBits( code, iCodeSize );

	// print the pure 8bit literal if it's encoded via escape code
	if ( m_Leafs[symbol] == -1 )
	{
		b &= writer.WriteBits( symbol, 8 );
		AddNewNode( symbol );
	}

	_ASSERT( b );

	return TRUE; // OK
}

//
// Purpose:
//   Reads bits form the ByteReader class example and traces parallel
//   the tree until a leaf node is reached
//
// Returns:
//   -1  - error
//
int AHuff::DecodeSymbol( ByteReader &reader )
{
	int    ret = -1;
	UINT   i   = ROOT_NODE;

	//
	// follow the bits to the final leaf
	//
    while ( !m_Nodes[i].bIsLeaf )
	{
        i  = m_Nodes[i].child;
        i += reader.ReadBit();
    }

	// grab literal
    ret = (int)m_Nodes[i].child;

	// handle an escape code
    if ( ret == ESCAPE )
	{
        ret = (BYTE)reader.ReadBits( 8 );
        AddNewNode( ret );
    }

	return (int)ret;
}